# DuckDB相关操作模块

import duckdb
import pandas as pd
import os
from dotenv import load_dotenv
from db.logger import logging

load_dotenv()

class DuckDBHelper:
    def __init__(self, db_path=None):
        # 优先从环境变量读取路径
        self.db_path = db_path or os.getenv('DUCKDB_PATH', 'data/local.duckdb')
    
    def get_conn(self):
        """
        获取DuckDB连接
        :return: DuckDB连接对象
        """
        conn=None
        try:
            conn = duckdb.connect(self.db_path)
        except Exception as e:
            logging.error(f"连接DuckDB失败: {e}")
            # raise
        finally:
            return conn

    def query_df(self, sql):
        """
        执行查询
        :return: DataFrame 或 None
        """
        conn = duckdb.connect(self.db_path)
        df = conn.execute(sql).df()
        conn.close()
        if df.empty:
            return None
        return df
    
    def query_sql_first(self, sql):
        """
        执行查询并返回第一行数据
        :param sql: 查询语句
        :return: 第一行数据的DataFrame或None
        """
        conn = duckdb.connect(self.db_path)
        df = conn.execute(sql).df()
        conn.close()
        if df.empty:
            return None
        return df.iloc[0]

    def batch_insert_df(self, df, table_name, pk_col=None):
        """
        批量导入DataFrame到DuckDB指定表，避免重复，失败自动回滚。
        :param df: 要导入的DataFrame
        :param table_name: 目标表名
        :param pk_col: 主键字段名（如有，避免重复导入）
        :return: None
        """
        if df is None or df.empty:
            logging.warning("无数据需要插入")
            return

        # 处理时间字段
        for col, dtype in df.dtypes.items():
            if str(dtype) == "datetime64[ns]":
                df[col] = df[col].dt.strftime('%Y-%m-%d %H:%M:%S')

        if not pk_col or pk_col not in df.columns:
            logging.error(f"未指定主键或主键不存在，程序中止。pk_col:{pk_col}, df columns: {df.columns.tolist()}")
            raise ValueError(f"未指定主键或主键不存在，程序中止。pk_col:{pk_col}, df columns: {df.columns.tolist()}") 

        try:
            conn = self.get_conn()
            if not conn:
                print('连接数据库失败,程序中止。')
                return
                
            # 删除主键重复的记录

            conn.execute("BEGIN TRANSACTION;")

            exist_sql = f"SELECT {pk_col} FROM {table_name}"
            try:
                exist_keys = set(self.query_df(exist_sql)[pk_col].tolist())
            except Exception:
                exist_keys = set()

            for content in  exist_keys:
                delete_sql = f"DELETE FROM {table_name} WHERE {pk_col} = '{content}'"
                conn.execute(delete_sql)    
           
                
            conn.register('df_view', df)
            conn.execute(f'INSERT INTO {table_name} SELECT * FROM df_view;')
            conn.execute("COMMIT;")
            logging.info(f"成功导入 {len(df)} 条数据到 {table_name}")
        except Exception as e:
            conn.execute("ROLLBACK;")
            logging.error(f"导入失败，已回滚: {e}")
            raise
        finally:
            conn.close()


    def batch_insert_sql(self, sql):
        """
        执行批量插入来自SQL查询结果的数据，异常时自动回滚。
        :param sql: 插入语句（支持多条）
        :return: 影响的行数
        """
        conn = duckdb.connect(self.db_path)
        affected_rows = 0
        try:
            conn.execute("BEGIN TRANSACTION;")
            result = conn.execute(sql)
            # 尝试获取影响行数
            try:
                affected_rows = result.rowcount if hasattr(result, 'rowcount') else 0
            except Exception:
                affected_rows = 0
            conn.execute("COMMIT;")
            logging.info(f"批量插入成功，影响行数: {affected_rows}")
        except Exception as e:
            conn.execute("ROLLBACK;")
            logging.error(f"批量插入失败，已回滚: {e}")
            raise
        finally:
            conn.close()
        return affected_rows

    def execute_sql(self, sql):
        """
        执行任意SQL语句（非查询）
        :param sql: SQL语句
        :return: None
        """
        conn = self.get_conn()
        if not conn:
            print('连接数据库失败,程序中止。')
            return

        try:
            conn.execute("BEGIN TRANSACTION;")
            conn.execute(sql)
            conn.execute("COMMIT;")
            logging.info(f"成功执行SQL: {sql}")
        except Exception as e:
            conn.execute("ROLLBACK;")
            logging.error(f"执行SQL失败: {e}")
            raise
        finally:
            conn.close()